#coding=utf-8
import tensorflow as tf
import numpy as np
import os

class vector():
	def __init__(self, 
				 embedding_size=128,
				 head_user_list=None,
				 user_list=None,
				 head_model_dir=None,
				 model_ckpt=None,
				 train_dir=None):

		self.embedding_size = embedding_size
		self.head_user_list = head_user_list
		self.user_list = user_list
		self.head_model_dir = head_model_dir
		self.model_ckpt = model_ckpt
		self.train_dir=train_dir

		self.build_graph()
		self.init_op()
		
		if self.head_model_dir != None and os.path.exists(self.head_model_dir):
			self.saver_head.restore(self.sess, tf.train.latest_checkpoint(self.head_model_dir))
		else:
			raise Exception, 'head_model_dir not exist!'

		if self.model_ckpt != None and os.path.exists(self.model_ckpt+'.index'):
			self.saver.restore(self.sess, self.model_ckpt)
		else:
			print('not found model_ckpt, init it')
	
	def init_op(self):
		self.sess = tf.Session(graph=self.graph)
		self.sess.run(self.init)
		self.sess.run(self.head_table.init)
		self.sess.run(self.table.init)

	def build_graph(self):
		self.graph = tf.Graph()
		with self.graph.as_default():
			with tf.name_scope('head'):
				self.embeddings = tf.get_variable('embeddings', shape=[len(self.head_user_list)+1, self.embedding_size],
					initializer=tf.random_uniform_initializer(minval=-1.0, maxval=1.0))

				self.head_table = tf.contrib.lookup.index_table_from_tensor(mapping=tf.convert_to_tensor(self.head_user_list),
																			num_oov_buckets=0,
																			default_value=len(self.head_user_list))

			with tf.name_scope('user'):
				self.embeddings_user = tf.get_variable('embeddings_user', shape=[len(self.user_list)+1, self.embedding_size],
					initializer=tf.random_uniform_initializer(minval=-1.0, maxval=1.0))
		
				self.table = tf.contrib.lookup.index_table_from_tensor(mapping=tf.convert_to_tensor(self.user_list),
																	   num_oov_buckets=0,
																	   default_value=len(self.user_list))
			self.init = tf.global_variables_initializer()
			self.saver_head = tf.train.Saver({'embeddings': self.embeddings})
			self.saver =  tf.train.Saver({'embeddings_user': self.embeddings_user})

			# 获取用户uid和followlist
			self.uid, self.follow_list = self.input_fn(self.train_dir)

			# 计算followlist的平均值
			self.follows_idx = self.head_table.lookup(self.follow_list)
			self.follows_embedding = tf.nn.embedding_lookup(self.embeddings, tf.squeeze(self.follows_idx))
			self.follows_mean = tf.reduce_mean(self.follows_embedding, 0)
			
			# 获取用户在embeddings_user中的位置(128,1)
			self.uid_idx = self.table.lookup(self.uid)
			self.uid_idx = tf.fill([self.embedding_size, 1], tf.reshape(self.uid_idx, []))
			
			# 生成用于更新embeddings_user的indices(128,1)
			self.indices = tf.range(0, self.embedding_size, 1, dtype=tf.int64)
			self.indices = tf.expand_dims(self.indices, 1)
			self.indices = tf.concat([self.uid_idx, self.indices], 1)	# (128, 2)

			# 根据indices结构更新embeddings_user
			self.update = tf.scatter_nd_update(self.embeddings_user, self.indices, self.follows_mean)

	def decode_line(self, line):
		columns = tf.string_split([line], '\t')
		follow_list = tf.string_split(columns.values[1:], ',')
		return columns.values[0], follow_list.values

	def input_fn(self, files_path):
		files = tf.data.Dataset.list_files(files_path)
		dataset = files.apply(tf.contrib.data.parallel_interleave(tf.data.TextLineDataset, cycle_length=6))
		dataset = dataset.map(self.decode_line, num_parallel_calls=1)
		dataset = dataset.filter(lambda x, y: tf.size(y) > 0)
		dataset = dataset.batch(1)
		iterator = dataset.make_one_shot_iterator()
		return iterator.get_next()

	def start(self):
		print('------------------------head embedding--------------------------')
		print(self.sess.run(self.embeddings))
		print('--------------------------start update---------------------------')
		try:
			i = 1
			while True:
				self.sess.run(self.update)
				i = i + 1
				if i % 10000 == 0:
					print('update %d users now' % i)
		except tf.errors.OutOfRangeError:
			print('update end!')

	def save_model(self):
		self.saver.save(self.sess, self.model_ckpt)
		print('------------------------save end-------------------------------')

	def print_model(self):
		print('------------------------user embedding-------------------------')
		self.saver.restore(self.sess, self.model_ckpt)
		print(self.sess.run(self.embeddings_user))
